import datetime
import multiprocessing
import os
import shlex
import logging
import shutil
import subprocess
import threading

thread_lock = threading.Lock()
process_lock = multiprocessing.Lock()
LOG = logging.getLogger(__name__)


def get_gpu_free_mem(gpu_id=0):
    import pynvml
    pynvml.nvmlInit()
    gpu_number = pynvml.nvmlDeviceGetCount()
    if gpu_number < 1:
        raise Exception('not find nvidia gpu!')
    handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    used = int(mem_info.used / 1024 / 1024)
    total = int(mem_info.total / 1024 / 1024)
    free = int(mem_info.free / 1024 / 1024)
    print('gpu memory info used={}MiB, total={}MiB'.format(used, total))
    pynvml.nvmlShutdown()
    return free


def run_subprocess_check_output(cmd: str, timeout=60):
    args = shlex.split(cmd)
    try:
        p = subprocess.check_output(args, stderr=subprocess.STDOUT, timeout=timeout)
    except subprocess.TimeoutExpired as time_e:
        err_msg = 'extract timeout={}, cmd={}'.format(timeout, cmd)
        LOG.error(err_msg)
        LOG.error(str(time_e))
        return False, err_msg
    except subprocess.CalledProcessError as call_e:
        err_msg = 'extract failed, maybe GPU RAM is too small, cmd={}'.format(cmd)
        LOG.error(err_msg)
        LOG.error(str(call_e))
        return False, err_msg
    else:
        msg = 'extract success, output={}, cmd={}'.format(p, cmd)
        LOG.info(msg)
        return True, p


def get_now_day():
    return datetime.datetime.now().strftime("%Y-%m-%d")


def make_dirs_if_not_exist(path: str):
    path = os.path.abspath(path)
    with process_lock:
        with thread_lock:
            if not os.path.exists(path):
                os.makedirs(path)


def get_path_tree(path, tree_files=None, ignore_list=None):
    if tree_files is None:
        tree_files = []
    file_list = os.listdir(path)
    for filename in file_list:
        if ignore_list and filename in ignore_list:
            continue
        current_path = os.path.join(path, filename)
        if os.path.isdir(current_path):
            tree_files.append({filename: get_path_tree(current_path, ignore_list=ignore_list)})
        else:
            tree_files.append(filename)
    return tree_files


def get_file_name(filepath):
    file = os.path.basename(filepath)
    return file.split('.')[0]


def get_number_file_fullname(path, number, length=3):
    file_list = os.listdir(path)
    number = str(number).rjust(length, '0')
    for filename in file_list:
        if filename.startswith(number):
            return filename
    else:
        raise Exception('not find number={} in {}'.format(number, path))


def split_audio(input_audio, time_split: int, output_dir):
    filename = get_file_name(input_audio)
    output_file_dir = os.path.join(output_dir, filename)
    if os.path.exists(output_file_dir):
        shutil.rmtree(output_file_dir)
    os.makedirs(output_file_dir)
    from pydub import AudioSegment
    sound = AudioSegment.from_mp3(input_audio)
    zhi = len(sound) / 1000 / time_split
    number = int(zhi) if zhi - int(zhi) < 0.2 else int(zhi) + 1
    number = number or 1
    time_split_ms = len(sound) // number
    for i in range(number):
        if i + 1 == number:
            audio = sound[i * time_split_ms:]
        else:
            audio = sound[i * time_split_ms:(i + 1) * time_split_ms]
        output_filename = '{}-{:.2f}s.mp3'.format(str(i + 1).rjust(3, '0'), len(audio) / 1000)
        output_filepath = os.path.join(output_file_dir, output_filename)
        audio.export(output_filepath, format="mp3")
